import os
import yaml
import pandas as pd
from src.assistant.dataloader import download_mnist_to_local, load_mnist_from_internet, load_mnist_from_local
from src.assistant.process import sup_process
from src.assistant.plot import plot_loss_curve, plot_model_results_bar
from src.assistant.text2mod import text2mod
import logging
from datetime import datetime

logging.basicConfig(level=logging.INFO)

def ensure_dirs(config):
    if config.get('plot', True):
        os.makedirs('fig', exist_ok=True)
    if config.get('evaluate', True):
        os.makedirs('eval_result', exist_ok=True)

def load_config(path):
    with open(path, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)
    # 需要分割的字段
    for key in ['model', 'kernel', 'base_estimator']:
        if key in config and isinstance(config[key], str):
            # 分割并去除首尾空格
            config[key] = [item.strip() for item in config[key].split(',')]
    # 其余单独字段的空字符串也转为None
    for key in ['train_size', 'test_ratio']:
        if key in config and (config[key] is None or (isinstance(config[key], str) and config[key].strip() == '')):
            config[key] = None
    return config

def write_results_table(result_diclist, filename='eval_result/result.txt'):
    """
    将模型评估结果写入markdown表格
    """
    # 只保留需要的字段
    df = pd.DataFrame(result_diclist)
    df = df[['model', 'accuracy', 'f1', 'train_time']]
    df.columns = ['model', 'accuracy', 'f1_score', 'train_time(s)']
    # 转为markdown表格字符串
    table_str = df.to_markdown(index=False)
    # 写入文件
    with open(filename, 'w', encoding='utf-8') as f:
        f.write(table_str + '\n')

    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    logging.info(f"[{now}] 评估结果已写入 {filename}")



def main():
    # 清空旧输出
    if os.path.exists('eval_result/result.txt'):
        open('eval_result/result.txt', 'w').close()
        now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        logging.info(f"[{now}] 已清空旧的评估结果文件 eval_result/result.txt")
    
    # 读取配置
    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    logging.info(f"[{now}] 开始读取配置文件 config.yaml")

    # config = load_config('config_svm.yaml')
    config = load_config('config_ada.yaml')

    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    logging.info(f"[{now}] 配置文件读取完成")

    ensure_dirs(config)

    # 数据加载
    mnist_path = os.path.join('dataset', 'mnist_784.pkl')
    if not os.path.exists(mnist_path):
        now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        logging.info(f"[{now}] 未检测到本地MNIST数据集，开始下载")
        download_mnist_to_local()
        logging.info(f"[{now}] MNIST数据集下载完成")
    else:
        now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        logging.info(f"[{now}] 检测到本地已有MNIST数据集，直接加载")

    X_train, X_test, y_train, y_test = load_mnist_from_local()
    # X_train, X_test, y_train, y_test = load_mnist_from_internet()
    logging.info(f"[{now}] MNIST数据集加载完成")
   
    # 模型训练与评估
    logging.info(f"[{now}] 开始模型训练与评估")
    result_diclist = sup_process(
        models=text2mod(config),
        X_train=X_train,
        y_train=y_train,
        X_test=X_test,
        y_test=y_test,
        evaluate=config.get('evaluate', False),
        Loss_curve=config.get('loss_curve', False)
    )
    logging.info(f"[{now}] 所有模型训练与评估完成")

    # 将结果写入带时间戳的文件
    now_str = datetime.now().strftime('%Y%m%d_%H%M%S')
    result_filename = f'eval_result/result_{now_str}.txt'
    write_results_table(result_diclist, filename=result_filename)

    if config.get('plot'):
        logging.info(f"[{now}] 开始绘制图表")
        if config.get('loss_curve'):
            plot_loss_curve(result_diclist)
        plot_model_results_bar(result_diclist)
        logging.info(f"[{now}] 图表绘制完成")

if __name__ == '__main__':
    main()
